{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true,
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import re\n",
    "import json\n",
    "import string\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import layers\n",
    "from tokenizers import BertWordPieceTokenizer\n",
    "from transformers import BertTokenizer,TFBertModel,BertConfig\n",
    "\n",
    "max_len = 384\n",
    "configuration = BertConfig()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "outputs": [],
   "source": [
    "slow_tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')\n",
    "save_path = 'bert_base_uncased/'\n",
    "if not os.path.exists(save_path):\n",
    "    os.makedirs(save_path)\n",
    "slow_tokenizer.save_pretrained(save_path)\n",
    "\n",
    "tokenizer = BertWordPieceTokenizer('bert_base_uncased/vocab.txt',lowercase=True)\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": false
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "outputs": [],
   "source": [
    "train_data_url = \"https://rajpurkar.github.io/SQuAD-explorer/dataset/train-v1.1.json\"\n",
    "train_path = keras.utils.get_file(\"train.json\", train_data_url)\n",
    "eval_data_url = \"https://rajpurkar.github.io/SQuAD-explorer/dataset/dev-v1.1.json\"\n",
    "eval_path = keras.utils.get_file(\"eval.json\", eval_data_url)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": false
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "outputs": [
    {
     "name": "stdout",
     "text": [
      "87599 training points created.\n",
      "10570 evaluation points created.\n"
     ],
     "output_type": "stream"
    }
   ],
   "source": [
    "class SquadExample:\n",
    "    def __init__(self,question,context,start_char_idx,answer_text,all_answers):\n",
    "        self.question = question\n",
    "        self.context = context\n",
    "        self.start_char_idx = start_char_idx\n",
    "        self.answer_text = answer_text\n",
    "        self.all_answers = all_answers\n",
    "        self.skip = False\n",
    "    \n",
    "    def preprocess(self):\n",
    "        context = self.context\n",
    "        question = self.question\n",
    "        answer_text = self.answer_text\n",
    "        start_char_idx = self.start_char_idx\n",
    "        \n",
    "        context = \" \".join(str(context).split())\n",
    "        question = \" \".join(str(question).split())\n",
    "        answer = \" \".join(str(answer_text).split())\n",
    "        \n",
    "        end_char_idx = start_char_idx + len(answer)\n",
    "        if end_char_idx >= len(context):\n",
    "            self.skip = True\n",
    "            return \n",
    "        is_char_in_ans = [0] * len(context)\n",
    "        for idx in range(start_char_idx,end_char_idx):\n",
    "            is_char_in_ans[idx] = 1\n",
    "        \n",
    "        tokenized_context = tokenizer.encode(context)\n",
    "        ans_token_idx = []\n",
    "        for idx,(start,end) in enumerate(tokenized_context.offsets):\n",
    "            if sum(is_char_in_ans[start:end]) > 0:\n",
    "                ans_token_idx.append(idx)\n",
    "        \n",
    "        if len(ans_token_idx) == 0:\n",
    "            self.skip = True\n",
    "            return \n",
    "        \n",
    "        start_token_idx = ans_token_idx[0]\n",
    "        end_token_idx = ans_token_idx[-1]\n",
    "        \n",
    "        \n",
    "        tokenized_question = tokenizer.encode(question)\n",
    "        input_ids = tokenized_context.ids + tokenized_question.ids[1:]\n",
    "        token_type_ids = [0] * len(tokenized_context.ids) + [1]*len(\n",
    "            tokenized_question.ids[1:]\n",
    "        )\n",
    "        attention_mask = [1] * len(input_ids)\n",
    "        \n",
    "        padding_length = max_len - len(input_ids)\n",
    "        if padding_length > 0:\n",
    "            input_ids = input_ids + ([0] * padding_length)\n",
    "            attention_mask = attention_mask + ([0] * padding_length)\n",
    "            token_type_ids = token_type_ids + ([0] * padding_length)\n",
    "        elif padding_length < 0 :\n",
    "            self.skip = True\n",
    "            return \n",
    "        \n",
    "        self.input_ids = input_ids\n",
    "        self.token_type_ids = token_type_ids\n",
    "        self.attention_mask = attention_mask\n",
    "        self.start_token_idx = start_token_idx\n",
    "        self.end_token_idx = end_token_idx\n",
    "        self.context_token_to_char = tokenized_context.offsets\n",
    "\n",
    "with open(train_path,encoding='utf-8') as f:\n",
    "    raw_train_data = json.load(f)\n",
    "\n",
    "with open(eval_path,encoding='utf-8') as f:\n",
    "    raw_eval_data =json.load(f)\n",
    "\n",
    "def create_squad_examples(raw_data):\n",
    "    squad_examples = []\n",
    "    for item in raw_data['data']:\n",
    "        for para in item['paragraphs']:\n",
    "            context = para['context']\n",
    "            for qa in para['qas']:\n",
    "                question = qa['question']\n",
    "                answer_text = qa['answers'][0]['text']\n",
    "                all_answers = [_['text'] for _ in qa['answers']]\n",
    "                start_char_idx = qa['answers'][0]['answer_start']\n",
    "                squad_eg = SquadExample(\n",
    "                    question,context,start_char_idx,\n",
    "                    answer_text,all_answers\n",
    "                )\n",
    "                squad_eg.preprocess()\n",
    "                squad_examples.append(squad_eg)\n",
    "    return squad_examples\n",
    "\n",
    "def create_inputs_targets(squad_examples):\n",
    "    dataset_dict = {\n",
    "        'input_ids':[],\n",
    "        'token_type_ids':[],\n",
    "        'attention_mask':[],\n",
    "        'start_token_idx':[],\n",
    "        'end_token_idx':[],\n",
    "    }\n",
    "    \n",
    "    for item in squad_examples:\n",
    "        if item.skip == False:\n",
    "            for key in dataset_dict:\n",
    "                dataset_dict[key].append(getattr(item,key))\n",
    "    \n",
    "    for key in dataset_dict:\n",
    "        dataset_dict[key]= np.array(dataset_dict[key])\n",
    "    \n",
    "    x = [\n",
    "        dataset_dict['input_ids'],\n",
    "        dataset_dict['token_type_ids'],\n",
    "        dataset_dict['attention_mask']\n",
    "    ]\n",
    "    y = [dataset_dict['start_token_idx'],dataset_dict['end_token_idx']]\n",
    "    return x,y\n",
    "\n",
    "train_squad_examples = create_squad_examples(raw_train_data)\n",
    "x_train,y_train = create_inputs_targets(train_squad_examples)\n",
    "print(f\"{len(train_squad_examples)} training points created.\")\n",
    "\n",
    "eval_squad_examples = create_squad_examples(raw_eval_data)\n",
    "x_val,y_eval = create_inputs_targets(eval_squad_examples)\n",
    "print(f\"{len(eval_squad_examples)} evaluation points created.\")\n",
    "\n",
    "    \n",
    "        \n",
    "        \n",
    "        \n",
    "            \n",
    "        \n",
    "        \n",
    "        \n",
    "        "
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": false
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def create_model():\n",
    "    encoder = TFBertModel.from_pretrained('bert-base-uncased')\n",
    "    \n",
    "    input_ids = layers.Input(shape=(max_len,),dtype=tf.int32)\n",
    "    token_type_ids = layers.Input(shape=(max_len,),dtype=tf.int32)\n",
    "    attention_mask = layers.Input(shape=(max_len,),dtype=tf.int32)\n",
    "    embedding = encoder(\n",
    "        input_ids,token_type_ids = token_type_ids,attention_mask=attention_mask\n",
    "    )[0]\n",
    "    \n",
    "    start_logits = layers.Dense(1,name='start_logit',use_bias=False)(embedding)\n",
    "    start_logits = layers.Flatten()(start_logits)\n",
    "    \n",
    "    end_logits = layers.Dense(1,name='end_logit',use_bias=False)(embedding)\n",
    "    end_logits = layers.Flatten()(end_logits)\n",
    "    \n",
    "    start_probs = layers.Activation(keras.activations.softmax)(start_logits)\n",
    "    end_probs = layers.Activation(keras.activations.softmax)(end_logits)\n",
    "    \n",
    "    model = keras.Model(\n",
    "        inputs = [input_ids,token_type_ids,attention_mask],\n",
    "        outputs = [start_probs,end_probs]\n",
    "    )\n",
    "    \n",
    "    loss = keras.losses.SparseCategoricalCrossentropy(from_logits=False)\n",
    "    optimizer = keras.optimizers.Adam(lr=5e-5)\n",
    "    model.compile(optimizer=optimizer,loss=[loss,loss])\n",
    "    return model\n",
    "    \n",
    "    \n",
    "\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": true
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [
    {
     "name": "stderr",
     "text": [
      "Widget Javascript not detected.  It may not be installed or enabled properly.\n"
     ],
     "output_type": "stream"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2c9bb1b736a84ae68799a15ad5b99127"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "model = create_model()\n",
    "model.summary()\n",
    "\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": true
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "conda-env-keras-py",
   "language": "python",
   "display_name": "Python [conda env:keras]"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "source": [],
    "metadata": {
     "collapsed": false
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}